package com.leal.kafka

import com.alibaba.fastjson.{JSON, JSONObject}
import com.leal.sink.KafkaSink
import com.leal.util.{KafkaUtil, SparkUtil}
import org.apache.spark.SparkContext
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.sql.{Dataset, SparkSession}
import org.apache.spark.util.LongAccumulator

/**
 * @Classname bigdata
 * @Description Kafka Simple Producer
 * @Date 2023/2/28 17:21
 * @Created by leal
 */
object KafkaProducer {

  def main(args: Array[String]): Unit = {
    //0- args
    val bootstrapServers = "127.0.0.1:9092"
    val topic = "quickstart-events"

    val needKey = false
    /**
     * if needKey is true
     * keyField="field1,field2"
     */
    //1- get hive data
    val spark: SparkSession = SparkUtil.initSpark(enableHive = true)
    val context: SparkContext = spark.sparkContext
    val count: LongAccumulator = context.longAccumulator("count")

    //2 sead to kafka
    val sql = "select * from leal_test.kafka_test_data;"

    val kafkaProducer: Broadcast[KafkaSink[String, String]] = KafkaUtil.kafkaProducer(context, bootstrapServers)
    // 默认发送方式，不设置key
    val frame: Dataset[String] = SparkUtil.getDataFrameBySql(spark, sql).toJSON
    frame.show(5, truncate = false)

    frame.repartition(1)
      .foreachPartition(
        (row: Iterator[String]) => {
          if (row.nonEmpty) {
            if (needKey) {
              row.foreach((record: String) => {
                // 反解析拼接key
                val jsonObject: JSONObject = JSON.parse(record).asInstanceOf[JSONObject]
                kafkaProducer.value.send(topic, "data-" + jsonObject.getString("user_id"), record)
                count.add(1)
              })
            } else {
              row.foreach((record: String) => {
                kafkaProducer.value.send(topic, record)
                count.add(1)
              })
            }
          }
          else {
            println("result is null ")
          }
        }
      )

    // 设置key
    println("recode : " + count)
    //3 close
    spark.close()
  }
}
